#!/usr/bin/python3
# coding=UTF-8

import rospy
from detection_msgs.msg import Result
import os
import csv

# 定义回调函数，处理接收到的Result消息
def result_callback(data):
    row = data.Row
    num = data.Num
    good = data.Good
    bad = data.Bad
    # 检查文件是否存在，如果不存在就创建它
    file_path = '/home/tta/list_of_good2.csv'
    if not os.path.exists(file_path):
        with open(file_path, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            # 写入表头
            writer.writerow(["房间", "好人", "坏人"])
    # 防止提前记入数据
    if row == 0 and num == 0:
        return
    # 读取 CSV 文件并排序（如果文件存在）
    data = []
    if os.path.exists(file_path):
        with open(file_path, 'r') as csvfile:
            reader = csv.reader(csvfile)
            # 跳过表头
            next(reader)
            for row_num_str, good_num, bad_num in reader:
                row_num = int(row_num_str.split('-')[0])
                num_in_file = int(row_num_str.split('-')[1])
                data.append((row_num, num_in_file, int(good_num), int(bad_num)))
        # 使用内置的sorted函数排序
        sorted_data = sorted(data, key=lambda x: (x[0], x[1]))
        # 将排序后的数据写入CSV文件
        with open(file_path, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(["房间", "好人", "坏人"])
            for row_num, num_in_file, good_num, bad_num in sorted_data:
                writer.writerow([f"{row_num}-{num_in_file}", good_num, bad_num])
    # 检查新数据是否已存在，如果存在则跳过
    with open(file_path, 'r') as csvfile:
        reader = csv.reader(csvfile)
        # 跳过表头
        next(reader)  
        for row_num_str, good_num, bad_num in reader:
            try:
                row_num = int(row_num_str.split('-')[0])
                num_in_file = int(row_num_str.split('-')[1])
                if row_num == row and num_in_file == num:
                    return
            except ValueError:
                rospy.logerr("Invalid data format in file, skipping.")
    # 写入新数据
    with open(file_path, 'a', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow([f"{row}-{num}", good, bad])
# 主函数
if __name__ == '__main__':
    # 初始化节点
    rospy.init_node('ztcar_result')
    # 订阅节点
    rospy.Subscriber('result_car', Result, result_callback)
    rospy.Subscriber('result_uav', Result, result_callback)
    rospy.spin()
